//
// Created by limo on 2025/9/27.
//

#pragma once

#include <iostream>
#include <string>
#include <map>
#include <memory>
#include <chrono>

#include <boost/asio.hpp>
#include <boost/bind/bind.hpp>
#include <ikcp.h>

namespace net {
    namespace asio = boost::asio;
    using boost::asio::ip::udp;

    extern "C" {
        // KCP 会通过这个函数将内部数据发送到 UDP Socket
        int kcpOutput(const char *buf, int len, ikcpcb *kcp, void *user);
    }

    // Session ID 类型
    using ConvIdType = IUINT32; // KCP 中的会话 ID 类型 (Conversation ID)
    using KcpMessageCallback = std::function<void(ConvIdType, const std::string&, std::size_t)>;
    using ErrorCallback = std::function<void(const boost::system::error_code&)>;

    class KcpSession : public std::enable_shared_from_this<KcpSession> {
    public:
        // 构造函数：需要 KCP ID、io_context 和一个发送函数
        KcpSession(asio::io_context& ioContext,
                   const udp::endpoint& remoteEndpoint,
                   ConvIdType convId,
                   std::function<void(const char*, int, const udp::endpoint&)> udpSendFunction)
                : m_convId(convId),
                  m_remoteEndpoint(remoteEndpoint),
                  m_udpSendFunction(std::move(udpSendFunction)),
                  m_updateTimer(ioContext)
        {
            // 1. 创建 KCP 控制块
            m_kcpControlBlock = ikcp_create(convId, this);

            // 2. 设置 KCP 输出回调（绑定到 kcpOutput）
            m_kcpControlBlock->output = kcpOutput;

            // 3. 设置 KCP 模式 (例如：快速模式)
            ikcp_nodelay(m_kcpControlBlock, 1, 10, 2, 1); // (nodelay, interval, resend, ncw)
            ikcp_wndsize(m_kcpControlBlock, 128, 128);

        }

        ~KcpSession() {
            if (m_kcpControlBlock) {
                ikcp_release(m_kcpControlBlock);
                m_kcpControlBlock = nullptr;
            }
            // 确保定时器停止
            m_updateTimer.cancel();
        }

        void start() {
            startKcpUpdateTimer();
        }

        // --- 公有接口 ---

        // 接收原始 UDP 数据并输入到 KCP
        void input(const char* data, int size) {
            ikcp_input(m_kcpControlBlock, data, size);
            // 数据输入后立即尝试接收 KCP 内部缓冲区的数据
            receiveData();
        }

        // 发送 KCP 消息
        int send(const std::string& message) {
            return ikcp_send(m_kcpControlBlock, message.c_str(), message.length());
        }

        // 设置消息回调
        void setMessageCallback(KcpMessageCallback cb) { m_onMessage = std::move(cb); }

        // 获取会话ID
        ConvIdType getConvId() const { return m_convId; }

        // 获取端点
        const udp::endpoint& getRemoteEndpoint() const { return m_remoteEndpoint; }

    private:
        // 成员变量
        ConvIdType m_convId;
        udp::endpoint m_remoteEndpoint;
        ikcpcb* m_kcpControlBlock;

        asio::steady_timer m_updateTimer;

        // 发送函数：将 KCP 封装的数据发送到 UDP socket
        std::function<void(const char*, int, const udp::endpoint&)> m_udpSendFunction;

        // 回调
        KcpMessageCallback m_onMessage;

        // --- KCP 内部逻辑 ---

        // 周期性更新 KCP 状态 (模拟时钟滴答)
        void startKcpUpdateTimer() {
            // KCP 建议的更新间隔，这里使用 10ms
            const auto updateInterval = std::chrono::milliseconds(10);
            m_updateTimer.expires_after(updateInterval);

            m_updateTimer.async_wait(
                    boost::bind(&KcpSession::handleKcpUpdate, shared_from_this(),
                                asio::placeholders::error)
            );
        }

        void handleKcpUpdate(const boost::system::error_code& error) {
            if (!error) {
                // 1. 更新 KCP 内部状态
                ikcp_update(m_kcpControlBlock, iclock());

                // 2. 重新启动定时器
                startKcpUpdateTimer();

            } else if (error != asio::error::operation_aborted) {
                std::cerr << "[KCP Update] Timer error: " << error.message() << std::endl;
            }
        }

        // 检查并接收 KCP 内部缓冲区中的完整数据包
        void receiveData() {
            // 检查是否有完整的 KCP 包可读
            int len;
            while ((len = ikcp_peeksize(m_kcpControlBlock)) > 0) {
                std::string buffer(len, '\0');
                int ret = ikcp_recv(m_kcpControlBlock, buffer.data(), len);

                if (ret > 0 && m_onMessage) {
                    // 触发消息回调
                    m_onMessage(m_convId, buffer, ret);
                }
            }
        }

        // 简单的时钟函数 (ikcp 需要毫秒时间戳)
        IUINT32 iclock() const {
            return static_cast<IUINT32>(std::chrono::duration_cast<std::chrono::milliseconds>(
                    std::chrono::steady_clock::now().time_since_epoch()
            ).count());
        }

        // 友元函数 kcpOutput 需要访问 m_udpSendFunction
        friend int net::kcpOutput(const char *buf, int len, ikcpcb *kcp, void *user);
    };

    // --- KCP 输出回调的实现 ---
    // 这个函数是 KCP 库的“胶水”代码，它将 KCP 封装的数据扔给 UDP socket 发送
    int kcpOutput(const char *buf, int len, ikcpcb *kcp, void *user) {
        if (user) {
            KcpSession* session = static_cast<KcpSession*>(user);
            // 调用 KcpSession 内部存储的 UDP 发送函数
            session->m_udpSendFunction(buf, len, session->m_remoteEndpoint);
        }
        return 0;
    }

    class KcpServer {
    public:
        KcpServer(asio::io_context& ioContext, short listenPort)
                : m_socket(ioContext, udp::endpoint(udp::v4(), listenPort)),
                  m_ioContext(ioContext),
                  m_nextConvId(100) // KCP ID 从 100 开始
        {
            std::cout << "KcpServer started on port " << listenPort << std::endl;
            startReceive();
        }

        // --- 公有方法：设置回调函数 ---
        void setMessageCallback(KcpMessageCallback cb) { m_onMessage = std::move(cb); }

        std::shared_ptr<KcpSession> getSessionByConvId(ConvIdType convId) {
            auto it = m_convIdToSessionMap.find(convId);
            if (it != m_convIdToSessionMap.end()) {
                return it->second;
            }
            return nullptr;
        }

        void setErrorCallback(ErrorCallback cb)      { m_onError = std::move(cb); }

    private:
        // 成员变量
        udp::socket m_socket;
        udp::endpoint m_remoteEndpoint;
        std::array<char, 2048> m_recvBuffer; // KCP 包可能较大
        asio::io_context& m_ioContext;
        ConvIdType m_nextConvId;

        // KCP 会话管理：使用 UDP EndPoint 作为 Key
        std::map<udp::endpoint, std::shared_ptr<KcpSession>> m_activeSessions;

        //  按 ConvId 查找 (用于 Echo/业务逻辑)
        std::map<ConvIdType, std::shared_ptr<KcpSession>> m_convIdToSessionMap;

        // 回调
        KcpMessageCallback m_onMessage;
        ErrorCallback m_onError;


        // --- 核心异步 I/O 逻辑 (基于 UDP) ---

        // 启动 UDP 接收
        void startReceive() {
            m_socket.async_receive_from(
                    asio::buffer(m_recvBuffer),
                    m_remoteEndpoint,
                    boost::bind(&KcpServer::handleReceive, this, asio::placeholders::error, asio::placeholders::bytes_transferred)
            );
        }

        // 处理收到的 UDP 数据包
        void handleReceive(const boost::system::error_code& error, std::size_t bytesTransferred) {
            if (!error) {
                // 1. 尝试从数据包头部提取 KCP 的 Conversation ID (Conv ID)
                ConvIdType convId = 0;
                if (bytesTransferred >= sizeof(ConvIdType)) {
                    // KCP 包头的前 4 字节就是 Conv ID
                    std::memcpy(&convId, m_recvBuffer.data(), sizeof(ConvIdType));
                    // 注意：这里需要处理网络字节序到主机字节序的转换（如果需要）
                    convId = ntohl(convId); // 假设我们使用网络字节序
                }

                // 2. 查找或创建 KcpSession
                auto it = m_activeSessions.find(m_remoteEndpoint);
                std::shared_ptr<KcpSession> sessionPtr;

                if (it == m_activeSessions.end()) {
                    // 新连接：检查 Conv ID 是否为 0 或非法，这里简化为直接创建新会话
                    convId = m_nextConvId++; // 分配新的 Conv ID

                    // 构造一个发送函数，传递给 KcpSession
                    auto sendFunc = [this](const char* data, int len, const udp::endpoint& ep) {
                        this->udpSend(data, len, ep);
                    };

                    // 创建新会话
                    sessionPtr = std::make_shared<KcpSession>(
                            m_ioContext, m_remoteEndpoint, convId, std::move(sendFunc)
                    );
                    sessionPtr->start();
                    // 将服务器的消息回调设置给 Session
                    sessionPtr->setMessageCallback(m_onMessage);
                    // 插入 map
                    m_activeSessions[m_remoteEndpoint] = sessionPtr;

                    m_convIdToSessionMap[convId] = sessionPtr;

                    std::cout << "[KCP] New Session " << convId << " from "
                              << m_remoteEndpoint.address().to_string() << ":" << m_remoteEndpoint.port() << std::endl;
                } else {
                    sessionPtr = it->second;
                }

                // 3. 将收到的原始 UDP 数据输入到 KCP 状态机
                sessionPtr->input(m_recvBuffer.data(), bytesTransferred);

            } else {
                if (m_onError) m_onError(error);
            }

            // 4. 保持监听
            startReceive();
        }

        // KCP 输出回调最终调用的 UDP 裸发函数
        void udpSend(const char* data, int len, const udp::endpoint& targetEndpoint) {
            m_socket.async_send_to(
                    asio::buffer(data, len),
                    targetEndpoint,
                    [](const boost::system::error_code& error, std::size_t /*bytesTransferred*/) {
                        if (error) {
                            std::cerr << "[KCP] UDP Send Error: " << error.message() << std::endl;
                            // 注意：这里的错误处理通常比较简单，因为 KCP 自身会处理重传
                        }
                    }
            );
        }
    };

} // net
